//
// Created by zy on 22-6-7.
//
//
// Created by zy on 22-6-7.
//
#include <cassert>
#include <chrono>
#include <iostream>
#include "leveldb/db.h"
#include "leveldb/comparator.h"
#include "util.h"
#include "stats.h"
#include "learned_index.h"
#include <cstring>
#include "cxxopts.hpp"
#include <unistd.h>
#include <fstream>
#include "../db/version_set.h"
#include <cmath>
#include <random>

using namespace std::chrono;
using namespace leveldb;
using namespace adgMod;
using std::string;
using std::cerr;
using std::cout;
using std::endl;
using std::to_string;
using std::vector;
using std::map;
using std::ifstream;
using std::string;

int num_pairs_base = 1000;
int mix_base = 20;
#include <set>


class NumericalComparator : public Comparator {
 public:
  NumericalComparator() = default;
  virtual const char* Name() const {return "adgMod:NumericalComparator";}
  virtual int Compare(const Slice& a, const Slice& b) const {
    uint64_t ia = adgMod::ExtractInteger(a.data(), a.size());
    uint64_t ib = adgMod::ExtractInteger(b.data(), b.size());
    if (ia < ib) return -1;
    else if (ia == ib) return 0;
    else return 1;
  }
  virtual void FindShortestSeparator(std::string* start, const Slice& limit) const { return; };
  virtual void FindShortSuccessor(std::string* key) const { return; };
};


void PutAndPrefetch(int lower, int higher, vector<string>& keys) {
  adgMod::Stats* instance = adgMod::Stats::GetInstance();

  Status status;

  instance->StartTimer(9);
  for (int i = lower; i < higher; ++i) {
    string value = generate_value(0);
    status = db->Put(write_options, keys[i], value);
    assert(status.ok() && "File Put Error");
  }
  instance->PauseTimer(9, true);

  //cout << "Put Complete" << endl;


  instance->StartTimer(10);
  for (int i = lower; i < higher; ++i) {
    string value;
    status = db->Get(read_options, keys[i], &value);
    //cout << "Get " << keys[i] << " Done" << endl;
    assert(status.ok() && "File Get Error");
  }
  instance->PauseTimer(10, true);

  //cout << "Prefetch Complete" << endl;
};

enum LoadType {
  Ordered = 0,
  Reversed = 1,
  ReversedChunk = 2,
  Random = 3,
  RandomChunk = 4
};


int main(int argc, char *argv[]) {
  int num_operations, num_iteration, num_mix;
  float test_num_segments_base;
  float num_pair_step;
  string db_location, profiler_out, input_filename, distribution_filename,
      ycsb_filename;
  bool print_single_timing, print_file_info, evict, unlimit_fd,
      use_distribution = false, pause, use_ycsb = false;
  bool change_level_load, change_file_load, change_level_learning,
      change_file_learning;
  int load_type, insert_bound, length_range;
  string db_location_copy;

  cxxopts::Options commandline_options("leveldb read test",
                                       "Testing leveldb read performance.");
  commandline_options.add_options()(
      "n,get_number", "the number of gets (to be multiplied by 1024)",
      cxxopts::value<int>(num_operations)->default_value("1000"))(
      "s,step", "the step of the loop of the size of db",
      cxxopts::value<float>(num_pair_step)->default_value("1"))(
      "i,iteration", "the number of iterations of a same size",
      cxxopts::value<int>(num_iteration)->default_value("1"))(
      "m,modification", "if set, run our modified version",
      cxxopts::value<int>(adgMod::MOD)->default_value("0"))(
      "h,help", "print help message",
      cxxopts::value<bool>()->default_value("false"))(
      "d,directory", "the directory of db",
      cxxopts::value<string>(db_location)->default_value("/mnt/ssd/testdb"))(
      "k,key_size", "the size of key",
      cxxopts::value<int>(adgMod::key_size)->default_value("8"))(
      "v,value_size", "the size of value",
      cxxopts::value<int>(adgMod::value_size)->default_value("8"))(
      "single_timing", "print the time of every single get",
      cxxopts::value<bool>(print_single_timing)->default_value("false"))(
      "file_info", "print the file structure info",
      cxxopts::value<bool>(print_file_info)->default_value("false"))(
      "test_num_segments", "test: number of segments per level",
      cxxopts::value<float>(test_num_segments_base)->default_value("1"))(
      "string_mode", "test: use string or int in model",
      cxxopts::value<bool>(adgMod::string_mode)->default_value("false"))(
      "e,model_error", "error in modesl",
      cxxopts::value<uint32_t>(adgMod::model_error)->default_value("8"))(
      "f,input_file", "the filename of input file",
      cxxopts::value<string>(input_filename)->default_value(""))(
      "multiple", "test: use larger keys",
      cxxopts::value<uint64_t>(adgMod::key_multiple)->default_value("1"))(
      "w,write", "writedb",
      cxxopts::value<bool>(fresh_write)->default_value("false"))(
      "c,uncache", "evict cache",
      cxxopts::value<bool>(evict)->default_value("false"))(
      "u,unlimit_fd", "unlimit fd",
      cxxopts::value<bool>(unlimit_fd)->default_value("false"))("x,dummy",
                                                                "dummy option")(
      "l,load_type", "load type",
      cxxopts::value<int>(load_type)->default_value("0"))(
      "filter", "use filter",
      cxxopts::value<bool>(adgMod::use_filter)->default_value("false"))(
      "mix", "mix read and write",
      cxxopts::value<int>(num_mix)->default_value("0"))(
      "distribution", "operation distribution",
      cxxopts::value<string>(distribution_filename)->default_value(""))(
      "change_level_load", "load level model",
      cxxopts::value<bool>(change_level_load)->default_value("false"))(
      "change_file_load", "enable level learning",
      cxxopts::value<bool>(change_file_load)->default_value("false"))(
      "change_level_learning", "load file model",
      cxxopts::value<bool>(change_level_learning)->default_value("false"))(
      "change_file_learning", "enable file learning",
      cxxopts::value<bool>(change_file_learning)->default_value("false"))(
      "p,pause", "pause between operation",
      cxxopts::value<bool>(pause)->default_value("false"))(
      "policy", "learn policy",
      cxxopts::value<int>(adgMod::policy)->default_value("0"))(
      "YCSB", "use YCSB trace",
      cxxopts::value<string>(ycsb_filename)->default_value(""))(
      "insert", "insert new value",
      cxxopts::value<int>(insert_bound)->default_value("0"))(
      "range", "use range query and specify length",
      cxxopts::value<int>(length_range)->default_value("0"));
  auto result = commandline_options.parse(argc, argv);
  if (result.count("help")) {
    printf("%s", commandline_options.help().c_str());
    exit(0);
  }

  std::default_random_engine e1(0), e2(255);
  srand(0);
  num_operations *= num_pairs_base;  // 1000 * 1000,什么意思呢？

  adgMod::fd_limit = unlimit_fd ? 1024 * 1024 : 1024;  // TODO: 需要测试一下是否真的影响性能
  adgMod::restart_read = true;
  adgMod::level_learning_enabled ^= change_level_learning;
  adgMod::file_learning_enabled ^= change_file_learning;
  adgMod::load_level_model ^= change_level_load;
  adgMod::load_file_model ^= change_file_load;

  // adgMod::file_learning_enabled = false;

  vector<uint64_t> distribution;  // uint64_t是unsigned long int的typedef别名
  vector<int> ycsb_is_write;
  // keys.reserve(100000000000 / adgMod::value_size);

  std::cerr << "读取keys from " << input_filename << std::endl;
  ifstream input(input_filename);
  assert(input.good());  // 加个判定，提前结束错误的情况
  std::string dev_id;
  char op_flag;
  std::string key;
  uint32_t op_len;
  uint64_t timestamp;
  vector<string> keys_all;
  vector<char> ops_all;
  while (input >> dev_id >> op_flag >> key >> op_len >> timestamp) {
    string the_key = generate_key(key);
    keys_all.push_back(std::move(the_key));
    ops_all.push_back(std::move(op_flag));
  }

  int N = int(keys_all.size() * 3.0/5.0);

  std::vector<std::string> keys_former(keys_all.begin(), keys_all.begin() + N);
  std::vector<std::string> keys_latter(keys_all.begin() + N, keys_all.end());
  std::vector<char> ops_former(ops_all.begin(), ops_all.begin() + N);
  std::vector<char> ops_latter(ops_all.begin() + N, ops_all.end());

  std::cerr << "keys_former:" << keys_former.size() << std::endl;
  std::cerr << "keys_latter:" << keys_latter.size() << std::endl;
  std::cerr << "ops_former:" << ops_former.size() << std::endl;
  std::cerr << "ops_latter:" << ops_latter.size() << std::endl;

  keys_all.clear();
  ops_all.clear();
  std::cerr << "读取并拆分keys完成" << std::endl;

  // 在load_ar的时候copy_out是false
  adgMod::Stats* instance = adgMod::Stats::GetInstance();
  vector<vector<size_t>> times(20);
  string values(1024 * 1024, '0');

  system("sync; echo 3 | sudo tee /proc/sys/vm/drop_caches");

  for (size_t iteration = 0; iteration < num_iteration; ++iteration) {
    std::uniform_int_distribution<uint64_t> uniform_dist_value(0, (uint64_t)values.size() - adgMod::value_size - 1);

    // 这里先把前半个数据集写入
    DB* db;
    Options options;
    ReadOptions& read_options = adgMod::read_options;
    WriteOptions& write_options = adgMod::write_options;
    Status status;
    options.create_if_missing = true;
    write_options.sync = false;
    instance->ResetAll();

    //    fresh_write要么为真要么为假，加载的时候不能跑测试，这样会导致learned index不对
    if (fresh_write && iteration == 0) {
      stage = "first";
      // 写入开始
      string command = "rm -rf " + db_location;
      system(command.c_str());
      system("sudo fstrim -a -v");
      std::cerr << ">>>>>>>>>>>>>>>>WRITING WRITING WRITING<<<<<<<<<<" << std::endl;

      status = DB::Open(options, db_location, &db);
      assert(status.ok() && "Open Error");

      instance->StartTimer(9);
      std::vector<std::string> keys = keys_former;
      assert(keys.size() == ops_former.size());
      int cut_size;
      // 对keys.size分情况讨论还是很有必要的，因为key的数量可能很少，不到10万。
      if (keys.size() > 100000) {
        cut_size = keys.size() / 100000;
      } else {
        cut_size = keys.size() / 1000;
      }
      std::cerr << "cut_size: " << cut_size << std::endl;
      std::vector<std::pair<int, int>> chunks;
      switch (load_type) {
        case Ordered: {
          for (int cut = 0; cut < cut_size; ++cut) {
            chunks.emplace_back(keys.size() * cut / cut_size,
                                keys.size() * (cut + 1) / cut_size);
          }
          break;
        }
        case ReversedChunk: {
          for (int cut = cut_size - 1; cut >= 0; --cut) {
            chunks.emplace_back(keys.size() * cut / cut_size,
                                keys.size() * (cut + 1) / cut_size);
          }
          break;
        }
        case Random: {
          std::random_shuffle(keys.begin(), keys.end());
          for (int cut = 0; cut < cut_size; ++cut) {
            chunks.emplace_back(keys.size() * cut / cut_size,
                                keys.size() * (cut + 1) / cut_size);
          }
          break;
        }
        case RandomChunk: {
          for (int cut = 0; cut < cut_size; ++cut) {
            chunks.emplace_back(keys.size() * cut / cut_size,
                                keys.size() * (cut + 1) / cut_size);
          }
          std::random_shuffle(chunks.begin(), chunks.end());
          break;
        }
        default:
          assert(false && "Unsupported load type");
      }

      // 加载数据库（放置数据，not ssd kind of），加载完了Bourbon才能学
      for (int cut = 0; cut < chunks.size(); ++cut) {
        std::cerr << "leveldb writing process: " << cut << " / " << chunks.size() << "\r";
        for (int i = chunks[cut].first; i < chunks[cut].second; ++i) {
          if (ops_former[i] == 'W') {
            status = db->Put(write_options, keys[i],
                             {values.data() + uniform_dist_value(e2),
                              (uint64_t)adgMod::value_size});
            assert(status.ok() && "File Put Error");
          } else if (ops_former[i] == 'R') {
            std::string value_load;
            status = db->Get(read_options, keys[i], &value_load);
            //            assert(status.ok() && "File Get Error");  // 这里不断言，因为不一定能找到，而且也不用打印not found
          } else throw "wrong operation flag";
        }
      }
      adgMod::db->vlog->Sync();
      instance->PauseTimer(9, true);
      std::cout << "Put Complete" << std::endl;
      // todo :上面加载的到底是leveldb还是wisckey?
      std::cerr << "数据库leveldb加载完成" << std::endl;
      keys.clear();

      if (print_file_info && iteration==0)  db->PrintFileInfo();
      adgMod::db->WaitForBackground();
      sleep(5);
      delete db;

      stage = "second";

      // 第二阶段，开始建立学习索引。注意建立学习索引是跟加载数据库一一对应的
      status = DB::Open(options, db_location, &db);
      adgMod::db->WaitForBackground();

      if (adgMod::MOD == 6 || adgMod::MOD == 7) {
        Version* current = adgMod::db->versions_->current();
        for (int i = 1; i < config::kNumLevels; ++i) {
          LearnedIndexData::Learn(new VersionAndSelf{current, adgMod::db->version_count,
                                                     current->learned_index_data_[i].get(), i});
        }
        current->FileLearn();
      }

      std::cout << "Shutting down" << std::endl;
      adgMod::db->WaitForBackground();
      delete db;
      std::cerr << "学习索引建立完成" << std::endl;
    }

    // 第三阶段，开始测试了，分为五个步骤
    if (!fresh_write) {
      // 加上drop cache呢
      readFileList((char*) db_location.data(), iteration);

      // 先划分keys
      int N_Period = keys_latter.size() / 5;
      std::vector<std::string> keys_latter_inPeriod(keys_latter.begin() + iteration * N_Period, keys_latter.begin() + (iteration+1) * N_Period);
      std::vector<char> ops_latter_inPeriod(ops_latter.begin() + iteration * N_Period, ops_latter.begin() + (iteration+1) * N_Period);

      stage = "second";
      std::cout << "Starting up" << std::endl;
      status = DB::Open(options, db_location, &db);
      adgMod::db->WaitForBackground();
      assert(status.ok() && "Open Error");

      uint64_t last_read = 0, last_write = 0;
      int last_level = 0, last_file = 0, last_baseline = 0, last_succeeded = 0,
          last_false = 0, last_compaction = 0, last_learn = 0;
      std::vector<uint64_t> detailed_times;
      bool start_new_event = true;


      int num_operations = keys_latter_inPeriod.size();
      std::cerr << "负载操作数写入../evaluation/num_ops.txt，为："
                << num_operations << std::endl;
      std::ofstream outfile("../evaluation/num_ops.txt");
      outfile << num_operations;
      outfile.close();
      std::cerr << "写入完成" << std::endl;

      instance->StartTimer(13);

      for (int i = 0; i < num_operations; ++i) {
        if (start_new_event) {
          detailed_times.push_back(instance->GetTime());
          start_new_event = false;
        }

        uint64_t index = i;
        const string& key = keys_latter_inPeriod[index];

        // all operations are Put
        instance->StartTimer(10);
        status =
            db->Put(write_options, key,
                    {values.data() + uniform_dist_value(e2), (uint64_t)64});
        instance->PauseTimer(10);
        assert(status.ok() && "Put Error");

        if ((i + 1) % (num_operations / 100) == 0) {
          detailed_times.push_back(instance->GetTime());
        }

        if ((i + 1) % (num_operations / 10) == 0) {
          int level_read = levelled_counters[0].Sum();
          int file_read = levelled_counters[1].Sum();
          int baseline_read = levelled_counters[2].Sum();
          int succeeded_read = levelled_counters[3].NumSum();
          int false_read = levelled_counters[4].NumSum();

          compaction_counter_mutex.Lock();
          int num_compaction = events[0].size();
          compaction_counter_mutex.Unlock();
          learn_counter_mutex.Lock();
          int num_learn = events[1].size();
          learn_counter_mutex.Unlock();

          uint64_t read_time = instance->ReportTime(4);
          uint64_t write_time = instance->ReportTime(10);
          std::pair<uint64_t, uint64_t> time = {detailed_times.front(),
                                                detailed_times.back()};

          events[2].push_back(new WorkloadEvent(
              time, level_read - last_level, file_read - last_file,
              baseline_read - last_baseline, succeeded_read - last_succeeded,
              false_read - last_false, num_compaction - last_compaction,
              num_learn - last_learn, read_time - last_read,
              write_time - last_write, std::move(detailed_times)));

          last_level = level_read;
          last_file = file_read;
          last_baseline = baseline_read;
          last_succeeded = succeeded_read;
          last_false = false_read;
          last_compaction = num_compaction;
          last_learn = num_learn;
          last_read = read_time;
          last_write = write_time;
          detailed_times.clear();
          start_new_event = true;
          std::cout << (i + 1) / (num_operations / 10) << std::endl;
          Version* current = adgMod::db->versions_->current();
          printf("LevelSize %d %d %d %d %d %d\n", current->NumFiles(0),
                 current->NumFiles(1), current->NumFiles(2),
                 current->NumFiles(3), current->NumFiles(4),
                 current->NumFiles(5));
        }
      }

      instance->PauseTimer(13, true);

      instance->ReportTime();
      for (int s = 0; s < times.size(); ++s) {
        times[s].push_back(instance->ReportTime(s));
      }

      adgMod::db->WaitForBackground();
      sleep(10);

      for (auto& event_array : events) {
        for (Event* e : event_array) e->Report();
      }

      for (Counter& c : levelled_counters) c.Report();

      file_data->Report();

      for (auto it : file_stats) {
        printf("FileStats %d %d %lu %lu %u %u %ld %d\n", it.first,
               it.second.level, it.second.start, it.second.end,
               it.second.num_lookup_pos, it.second.num_lookup_neg,
               it.second.size, it.first < file_data->watermark ? 0 : 1);
      }

      adgMod::learn_cb_model->Report();

      delete db;
      std::cerr << "iteration " << iteration << " finished........" << '\r';
    }
  }
  std::cerr << "All iterations finished~~~~~~~~~~~~~~~~~~~~~~~~~" << std::endl;

  // 统计部分，打印到文件里面
  if (!fresh_write) {
    perf_exclude_1st_Iteration(times);
  }

}